#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
数据库初始化脚本

用于创建MySQL数据库和表结构，支持从CSV文件导入历史数据

使用方法:
    python init_database.py --import-csv [CSV目录路径]
"""

import pymysql
import argparse
import os
import pandas as pd
import glob
from datetime import datetime
import sys

# 添加项目根目录到系统路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from config.db_config import DB_CONFIG
from scraper.data_saver import init_db, AirQualityRecord, save_data_to_mysql

def create_database(host, user, password, dbname, port=3306, charset='utf8mb4'):
    """
    创建MySQL数据库
    """
    try:
        # 创建连接
        connection = pymysql.connect(
            host=host,
            user=user,
            password=password,
            port=port,
            charset=charset
        )
        
        with connection.cursor() as cursor:
            # 创建数据库
            cursor.execute(f"CREATE DATABASE IF NOT EXISTS {dbname} CHARACTER SET {charset}")
            print(f"数据库 '{dbname}' 创建成功或已存在")
            
        connection.close()
        return True
    except Exception as e:
        print(f"创建数据库失败: {str(e)}")
        return False

def import_csv_data(csv_dir, session_maker):
    """
    从CSV文件导入历史数据到数据库
    
    参数:
        csv_dir: CSV文件目录
        session_maker: SQLAlchemy会话工厂
    """
    if not os.path.exists(csv_dir):
        print(f"目录不存在: {csv_dir}")
        return False
    
    # 获取所有CSV文件
    csv_files = glob.glob(os.path.join(csv_dir, "*.csv"))
    if not csv_files:
        print(f"未找到CSV文件: {csv_dir}")
        return False
    
    print(f"找到 {len(csv_files)} 个CSV文件")
    
    total_records = 0
    for csv_file in csv_files:
        try:
            # 读取CSV文件
            df = pd.read_csv(csv_file)
            print(f"处理文件: {os.path.basename(csv_file)}, {len(df)} 条记录")
            
            # 转换数据格式
            records = []
            for _, row in df.iterrows():
                # 尝试解析日期时间
                try:
                    if 'timestamp' in row:
                        timestamp = datetime.strptime(row['timestamp'], '%Y-%m-%d %H:%M:%S')
                    elif 'date' in row:
                        timestamp = datetime.strptime(row['date'], '%Y-%m-%d')
                    else:
                        # 使用文件名中的日期
                        filename = os.path.basename(csv_file)
                        date_str = filename.split('_')[-1].split('.')[0]  # 提取日期部分
                        timestamp = datetime.strptime(date_str, '%Y%m%d')
                except Exception as e:
                    print(f"日期解析失败，使用当前时间: {str(e)}")
                    timestamp = datetime.now()
                
                # 创建记录
                record = {
                    "timestamp": timestamp,
                    "PM2_5": float(row.get("PM2.5", 0)) if "PM2.5" in row else float(row.get("PM2_5", 0)),
                    "PM10": float(row.get("PM10", 0)),
                    "SO2": float(row.get("SO2", 0)),
                    "NO2": float(row.get("NO2", 0)),
                    "CO": float(row.get("CO", 0)),
                    "O3": float(row.get("O3", 0)),
                    "AQI": int(row.get("AQI", 0))
                }
                records.append(record)
            
            # 保存到数据库
            if records:
                success = save_data_to_mysql(records, session_maker)
                if success:
                    total_records += len(records)
                    print(f"成功导入 {len(records)} 条记录")
                else:
                    print(f"导入失败: {csv_file}")
        except Exception as e:
            print(f"处理文件 {csv_file} 失败: {str(e)}")
    
    print(f"总共导入 {total_records} 条记录")
    return total_records > 0

def main():
    parser = argparse.ArgumentParser(description='初始化空气质量数据库')
    parser.add_argument('--import-csv', dest='csv_dir', help='导入CSV文件目录')
    args = parser.parse_args()
    
    # 创建数据库
    print("开始创建数据库...")
    success = create_database(
        host=DB_CONFIG['host'],
        user=DB_CONFIG['user'],
        password=DB_CONFIG['password'],
        dbname=DB_CONFIG['database'],
        port=DB_CONFIG['port'],
        charset=DB_CONFIG['charset']
    )
    
    if not success:
        print("数据库创建失败，退出程序")
        return
    
    # 初始化表结构
    print("初始化数据库表结构...")
    try:
        session_maker = init_db()
        print("数据库表结构初始化成功")
    except Exception as e:
        print(f"数据库表结构初始化失败: {str(e)}")
        return
    
    # 导入CSV数据
    if args.csv_dir:
        print(f"开始从 {args.csv_dir} 导入CSV数据...")
        import_csv_data(args.csv_dir, session_maker)
    
    print("数据库初始化完成")

if __name__ == "__main__":
    main()